import requests
import json
import datetime
import time
import logging
import openpyxl
import re
from openpyxl.styles import Font, PatternFill, Alignment
from openpyxl.utils import get_column_letter
from typing import Dict, List, Optional

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

# 服务列表及对应的CVE搜索关键词
SERVICES = {
    "RabbitMQ": "RabbitMQ",
    "Nginx": "Nginx",
    "Redis": "Redis",
    "Zookeeper": "Zookeeper",
    "Nacos": "Nacos",
    "GitLab": "GitLab",
    "Kafka": "Apache Kafka",
    "MySQL": "MySQL",
    "Linux Kernel": "Linux kernel",  # 添加Linux内核监控
    "CentOS Kernel": "CentOS Linux kernel",  # 添加CentOS内核监控
    "Ubuntu Kernel": "Ubuntu Linux kernel"  # 添加Ubuntu内核监控
}

# NVD API配置
NVD_API_URL = "https://services.nvd.nist.gov/rest/json/cves/2.0"
REQUEST_DELAY = 6  # NVD API要求每分钟不超过5次请求（12秒/次），这里设置为6秒确保安全

# 企业微信机器人配置
# key= 修改为自己企业微信机器人值
WECOM_WEBHOOK = "https://qyapi.weixin.qq.com/cgi-bin/webhook/send?key="

# 颜色定义
SEVERITY_COLORS = {
    "CRITICAL": "FFFF0000",  # 红色
    "HIGH": "FFFFA500",  # 橙色
    "MEDIUM": "FFFFFF00",  # 黄色
    "LOW": "FF90EE90",  # 浅绿色
    "UNKNOWN": "FFD3D3D3"  # 灰色
}

# 发行版关键词映射
DISTRO_KEYWORDS = {
    "centos": ["centos", "rhel", "red hat enterprise linux"],
    "ubuntu": ["ubuntu", "canonical"]
}


def extract_distro_info(description: str) -> str:
    """
    从漏洞描述中提取影响的发行版信息
    """
    distros = set()

    # 检查描述中是否包含发行版关键词
    description_lower = description.lower()
    for distro, keywords in DISTRO_KEYWORDS.items():
        for keyword in keywords:
            if keyword in description_lower:
                distros.add(distro.capitalize())

    # 检查版本号模式
    version_patterns = [
        r"ubuntu[ -_]?(\d+\.\d+)",
        r"centos[ -_]?(\d+)",
        r"rhel[ -_]?(\d+)"
    ]

    versions = []
    for pattern in version_patterns:
        matches = re.findall(pattern, description_lower)
        if matches:
            versions.extend(matches)

    # 格式化发行版信息
    if distros:
        distro_info = ", ".join(sorted(distros))
        if versions:
            distro_info += f" ({', '.join(set(versions))})"
        return distro_info

    return ""


def get_recent_cves(service_name: str, days: int = 7) -> List[Dict]:
    """
    从NVD获取指定服务最近几天的CVE漏洞信息
    """
    end_date = datetime.datetime.now()
    start_date = end_date - datetime.timedelta(days=days)

    params = {
        "keywordSearch": service_name,
        "pubStartDate": start_date.strftime("%Y-%m-%dT%H:%M:%S.000"),
        "pubEndDate": end_date.strftime("%Y-%m-%dT%H:%M:%S.000"),
        "resultsPerPage": 20
    }

    try:
        time.sleep(REQUEST_DELAY)  # 遵守NVD API速率限制
        response = requests.get(NVD_API_URL, params=params)
        response.raise_for_status()
        data = response.json()

        cves = []
        for vuln in data.get("vulnerabilities", []):
            cve = vuln["cve"]
            metrics = cve.get("metrics", {})
            cvss_metric = next(iter(metrics.get("cvssMetricV31", metrics.get("cvssMetricV30", []))), None)

            # 提取发行版信息
            description = cve["descriptions"][0]["value"]
            distro_info = ""
            if "kernel" in service_name.lower():
                distro_info = extract_distro_info(description)

            cves.append({
                "id": cve["id"],
                "published": cve["published"],
                "description": description,
                "cvss_score": cvss_metric["cvssData"]["baseScore"] if cvss_metric else "N/A",
                "severity": cvss_metric["cvssData"]["baseSeverity"] if cvss_metric else "UNKNOWN",
                "url": f"https://nvd.nist.gov/vuln/detail/{cve['id']}",
                "distro": distro_info  # 添加发行版信息
            })
        return cves
    except Exception as e:
        logging.error(f"获取 {service_name} CVE失败: {str(e)}")
        return []


def format_message(service_cves: Dict[str, List[Dict]]) -> str:
    """
    格式化企业微信机器人消息
    """
    if not any(service_cves.values()):
        return "🎉 所有服务在过去7天内均未发现新漏洞"

    message = "🔍 **最新安全漏洞警报**\n\n"
    for service, cves in service_cves.items():
        if not cves:
            continue

        # 添加内核漏洞的特殊标识
        kernel_icon = "🐧 " if "kernel" in service.lower() else "🛡 "
        message += f"### {kernel_icon}{service}\n"

        for cve in cves:
            # 截断长描述
            description = cve['description']
            if len(description) > 200:
                description = description[:200] + "..."

            # 添加发行版信息
            distro_info = ""
            if cve.get('distro'):
                distro_info = f"\n  **影响系统**: {cve['distro']}"

            message += (
                f"- [{cve['id']}]({cve['url']}) "
                f"({cve['published'][:10]})\n"
                f"  **严重性**: {cve['severity']} ({cve['cvss_score']})"
                f"{distro_info}\n"
                f"  **描述**: {description}\n\n"
            )
        message += "\n"

    message += "> ℹ️ 数据来源: NVD国家漏洞数据库\n"
    message += "> 🐧 内核漏洞已标记影响系统信息"
    return message


def send_to_wecom(message: str):
    """
    发送消息到企业微信机器人
    """
    headers = {"Content-Type": "application/json"}
    payload = {
        "msgtype": "markdown",
        "markdown": {
            "content": message
        }
    }

    try:
        response = requests.post(WECOM_WEBHOOK, headers=headers, data=json.dumps(payload))
        response.raise_for_status()
        logging.info("消息发送成功")
    except Exception as e:
        logging.error(f"发送消息失败: {str(e)}")


def save_to_excel(service_cves: Dict[str, List[Dict]]):
    """
    将CVE数据保存到Excel文件
    """
    # 创建Excel工作簿
    wb = openpyxl.Workbook()
    wb.remove(wb.active)  # 移除默认工作表

    # 创建汇总工作表
    summary_sheet = wb.create_sheet(title="漏洞汇总")

    # 设置汇总表头
    summary_headers = ["服务名称", "漏洞数量", "最高严重性", "最后更新时间"]
    summary_sheet.append(summary_headers)

    # 样式设置
    header_font = Font(bold=True, color="FFFFFF")
    header_fill = PatternFill(start_color="1F4E78", end_color="1F4E78", fill_type="solid")
    center_align = Alignment(horizontal="center", vertical="center")

    # 应用表头样式
    for col in range(1, len(summary_headers) + 1):
        cell = summary_sheet.cell(row=1, column=col)
        cell.font = header_font
        cell.fill = header_fill
        cell.alignment = center_align

    # 处理每个服务的CVE数据
    for service, cves in service_cves.items():
        # 创建服务工作表
        sheet_name = service[:31]  # 工作表名称最多31字符
        sheet = wb.create_sheet(title=sheet_name)

        # 设置工作表头
        headers = ["CVE ID", "发布日期", "描述", "CVSS分数", "严重性", "影响系统", "详情链接"]
        sheet.append(headers)

        # 应用表头样式
        for col in range(1, len(headers) + 1):
            cell = sheet.cell(row=1, column=col)
            cell.font = header_font
            cell.fill = header_fill
            cell.alignment = center_align

        # 添加数据行
        if not cves:
            sheet.append(["无漏洞"])
        else:
            for cve in cves:
                row = [
                    cve['id'],
                    cve['published'][:10],
                    cve['description'],
                    cve['cvss_score'],
                    cve['severity'],
                    cve.get('distro', ''),
                    f'=HYPERLINK("{cve["url"]}", "查看详情")'
                ]
                sheet.append(row)

        # 设置列宽和自动换行
        for col in sheet.columns:
            max_length = 0
            column = col[0].column_letter
            for cell in col:
                try:
                    if len(str(cell.value)) > max_length:
                        max_length = len(str(cell.value))
                except:
                    pass
            adjusted_width = (max_length + 2) * 1.2
            sheet.column_dimensions[column].width = min(adjusted_width, 50)

        # 设置描述列自动换行
        for row in sheet.iter_rows(min_row=2, min_col=3, max_col=3):
            for cell in row:
                cell.alignment = Alignment(wrap_text=True)

        # 应用严重性颜色
        for row in sheet.iter_rows(min_row=2, min_col=5, max_col=5):
            for cell in row:
                severity = str(cell.value).upper()
                color = SEVERITY_COLORS.get(severity, SEVERITY_COLORS["UNKNOWN"])
                cell.fill = PatternFill(start_color=color, end_color=color, fill_type="solid")

        # 更新汇总表
        severity_levels = ["CRITICAL", "HIGH", "MEDIUM", "LOW", "UNKNOWN"]
        max_severity = "无漏洞"

        if cves:
            severities = [cve['severity'].upper() for cve in cves]
            for level in severity_levels:
                if level in severities:
                    max_severity = level
                    break

        summary_sheet.append([
            service,
            len(cves),
            max_severity,
            datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        ])

    # 设置汇总表样式
    for row in summary_sheet.iter_rows(min_row=2):
        # 设置漏洞数量颜色
        if row[1].value > 0:
            row[1].font = Font(color="FF0000", bold=True)  # 红色

        # 设置严重性颜色
        severity = str(row[2].value).upper()
        if severity != "无漏洞":
            color = SEVERITY_COLORS.get(severity, SEVERITY_COLORS["UNKNOWN"])
            row[2].fill = PatternFill(start_color=color, end_color=color, fill_type="solid")

    # 调整汇总表列宽
    for col in range(1, len(summary_headers) + 1):
        column_letter = get_column_letter(col)
        summary_sheet.column_dimensions[column_letter].width = 20

    # 保存文件
    filename = f"CVE_Report_{datetime.datetime.now().strftime('%Y%m%d_%H%M%S')}.xlsx"
    wb.save(filename)
    logging.info(f"报告已保存至: {filename}")
    return filename


def main():
    """
    主函数：获取所有服务的漏洞并发送通知
    """
    service_cves = {}
    for service, keyword in SERVICES.items():
        logging.info(f"正在获取 {service} 的漏洞信息...")
        cves = get_recent_cves(keyword)
        service_cves[service] = cves

    # 发送企业微信通知
    message = format_message(service_cves)
    send_to_wecom(message)

    # 保存Excel报告
    save_to_excel(service_cves)


if __name__ == "__main__":
    main()